'''
https://leetcode.cn/problems/trapping-rain-water-ii/description/
'''
import heapq
from typing import List

import numpy as np


class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        m, n = len(heightMap), len(heightMap[0])
        heap = []
        visited = np.full((m, n), False, dtype=bool)
        visited[:, 0] = True
        visited[:, -1] = True
        visited[0, :] = True
        visited[-1, :] = True
        # 4周堆： 堆：（z,x,y）
        for i in range(m):
            heapq.heappush(heap, (heightMap[i][0], i, 0))
            heapq.heappush(heap, (heightMap[i][n - 1], i, n - 1))
        for j in range(1, n - 1):
            heapq.heappush(heap, (heightMap[0][j], 0, j))
            heapq.heappush(heap, (heightMap[m - 1][j], m - 1, j))

        ans = 0
        while heap:
            # 每次先弹出最薄弱的地方，
            z, x, y = heapq.heappop(heap)  # z是高度，自己的高度 if 足够高 else 之前位置能帮忙撑起来的高度
            for nx, ny in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]:
                if 0 <= nx < m and 0 <= ny < n and not visited[nx, ny]:
                    heapq.heappush(heap, (max(heightMap[nx][ny], z), nx, ny))
                    ans += max(0, z - heightMap[nx][ny])
                    visited[nx, ny] = True
        return ans
